{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import string\n",
    "def normalize(s):\n",
    "    def remove_articles(text):\n",
    "        return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\n",
    "\n",
    "    def white_space_fix(text):\n",
    "        return \" \".join(text.split())\n",
    "\n",
    "    def remove_punc(text):\n",
    "        exclude = set(string.punctuation)\n",
    "        return \"\".join(ch for ch in text if ch not in exclude)\n",
    "\n",
    "    def lower(text):\n",
    "        return text.lower()\n",
    "\n",
    "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
    "\n",
    "def compute_f1(prediction, ground_truth):\n",
    "    if prediction is None:\n",
    "        return 0.0\n",
    "    prediction_tokens = normalize(prediction).split()\n",
    "    ground_truth_tokens = normalize(ground_truth).split()\n",
    "\n",
    "    common = set(prediction_tokens) & set(ground_truth_tokens)\n",
    "    num_same = len(common)\n",
    "\n",
    "    if num_same == 0:\n",
    "        return 0.0\n",
    "\n",
    "    precision = num_same / len(prediction_tokens)\n",
    "    recall = num_same / len(ground_truth_tokens)\n",
    "    f1 = 2 * precision * recall / (precision + recall)\n",
    "    return f1\n",
    "\n",
    "def exact_match_score(prediction, ground_truth):\n",
    "    if prediction is None:\n",
    "        return 0.0\n",
    "    return int(normalize(prediction) == normalize(ground_truth))\n",
    "\n",
    "def evaluate(predictions):\n",
    "    total = len(predictions)\n",
    "    f1_total = 0\n",
    "    em_total = 0\n",
    "\n",
    "    for item in predictions:\n",
    "        # if item['pred_answer_ori'] == None:\n",
    "        pred = item['pred_answer']\n",
    "        # else:\n",
    "        #     pred = item['pred_answer_ori']\n",
    "        gts = item['gt']\n",
    "\n",
    "        # 若gt是str，统一转换为列表处理\n",
    "        if isinstance(gts, str):\n",
    "            gts = [gts]\n",
    "\n",
    "        f1 = max([compute_f1(pred, gt) for gt in gts])\n",
    "        em = max([exact_match_score(pred, gt) for gt in gts])\n",
    "        if em == 1:\n",
    "            f1 = 1\n",
    "\n",
    "        f1_total += f1\n",
    "        em_total += em\n",
    "\n",
    "    return {\n",
    "        \"avg_f1\": f1_total / total if total > 0 else 0,\n",
    "        \"avg_em\": em_total / total if total > 0 else 0\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# results list：\n",
    "json_file_paths = [\n",
    "    \"7b_step200_musique_web_n4_0.json\",\n",
    "    \"7b_step200_musique_web_n4_1.json\",\n",
    "    \"7b_step200_musique_web_n4_2.json\",\n",
    "    \"7b_step200_musique_web_n4_3.json\",\n",
    "    \"7b_step200_musique_web_n4_4.json\",\n",
    "    \"7b_step200_musique_web_n4_5.json\",\n",
    "    \"7b_step200_musique_web_n4_6.json\",\n",
    "    \"7b_step200_musique_web_n4_7.json\",\n",
    "]\n",
    "\n",
    "# combined to one file\n",
    "merged_data = []\n",
    "for path in json_file_paths:\n",
    "    with open(path, \"r\", encoding=\"utf-8\") as f:\n",
    "        data = json.load(f)\n",
    "        merged_data.extend(data)\n",
    "\n",
    "print(len(merged_data))\n",
    "# save\n",
    "with open(\"7b_step200_musique_web_n4.json\", \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(merged_data, f, ensure_ascii=False, indent=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('./evaluation_results/musique/7b_step200_musique_web_n4.json', 'r') as f:\n",
    "    combine_results = json.load(f)\n",
    "print(len(combine_results))\n",
    "\n",
    "count_none = 0\n",
    "for item in combine_results:\n",
    "    if item['pred_answer'] == None:\n",
    "        count_none += 1\n",
    "print(count_none)\n",
    "results = evaluate(combine_results)\n",
    "print(\"Average F1:\", results['avg_f1'])\n",
    "print(\"Average EM:\", results['avg_em'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "r1-v",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
